"""
分类服务主文件

基于FastAPI实现的图像分类服务，支持通过环境变量动态配置模型和端口
"""

import os
import sys
import io
import base64
import numpy as np
import logging
from typing import Optional, List
from fastapi import FastAPI, HTTPException, UploadFile, File, Request
from fastapi.responses import JSONResponse
from fastapi.exceptions import RequestValidationError
from pydantic import BaseModel
from PIL import Image
from pathlib import Path

# 添加项目根目录到Python路径
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(current_dir)
sys.path.insert(0, project_root)

# 添加项目根目录到Python路径
project_root = str(Path(os.path.abspath(__file__)).parent.parent.parent)
sys.path.insert(0, project_root)

# 导入分类器和分割器
from src.JinhengLib import Classifier, Segmenter

# 从环境变量读取配置（必须通过环境变量配置）
MODEL_PATH = os.getenv('MODEL_PATH')  # 分类模型路径（可选）
SEGMENT_MODEL_PATH = os.getenv('SEGMENT_MODEL_PATH')  # 分割模型路径（可选）
PORT = int(os.getenv('PORT', 8796))
HOST = os.getenv('HOST', '0.0.0.0')

# GPU配置
USE_GPU = os.getenv('USE_GPU', 'false').lower() == 'true'  # 是否使用GPU
GPU_DEVICE_IDS = os.getenv('GPU_DEVICE_IDS', '[]')  # GPU设备ID列表（JSON格式字符串）
CUDA_VISIBLE_DEVICES = os.getenv('CUDA_VISIBLE_DEVICES', '')  # CUDA可见设备（Docker配置）


def get_device() -> str:
    """
    根据配置和实际环境动态获取推理设备
    
    重要说明（多GPU场景）：
    - Docker的device_ids指定宿主机GPU编号，容器内会重新映射为从0开始编号
    - CUDA_VISIBLE_DEVICES应该设置为容器内的GPU编号（从0开始），而不是宿主机编号
    - 例如：device_ids: ["1"] → 容器内只有cuda:0 → CUDA_VISIBLE_DEVICES=0
    - 例如：device_ids: ["0","1"] → 容器内有cuda:0,cuda:1 → CUDA_VISIBLE_DEVICES=0,1
    - 因此，如果设置了 CUDA_VISIBLE_DEVICES，应该使用 cuda:0（容器内第一个可见GPU）
    
    优先级：
    1. 如果配置了 USE_GPU=true 且 CUDA 可用，使用 cuda
    2. 否则使用 cpu
    
    Returns:
        str: 设备名称 ('cuda:0' 或 'cpu')
    """
    # 使用根logger（与setup_logging()配置的logger一致）
    logger = logging.getLogger()
    
    if USE_GPU:
        logger.info(f"GPU配置: USE_GPU=true, GPU_DEVICE_IDS={GPU_DEVICE_IDS}, CUDA_VISIBLE_DEVICES={CUDA_VISIBLE_DEVICES}")
        try:
            import torch
            if torch.cuda.is_available():
                visible_count = torch.cuda.device_count()
                logger.info(f"CUDA可用，容器内检测到 {visible_count} 个GPU设备")
                
                # 如果设置了 CUDA_VISIBLE_DEVICES，说明Docker已经限制了可见GPU
                # 此时PyTorch在容器内只能看到指定的GPU，并且会重新映射为从0开始
                # CUDA_VISIBLE_DEVICES应该已经是容器内的GPU编号（从0开始）
                # 因此应该使用 cuda:0（容器内第一个可见GPU）
                if CUDA_VISIBLE_DEVICES:
                    device = "cuda:0"
                    logger.info(f"使用GPU设备: {device} (CUDA_VISIBLE_DEVICES={CUDA_VISIBLE_DEVICES}，容器内GPU编号从0开始)")
                    return device
                
                # 如果没有设置 CUDA_VISIBLE_DEVICES，从 GPU_DEVICE_IDS 解析
                # 这种情况通常不会发生，因为Docker配置中会设置 CUDA_VISIBLE_DEVICES
                try:
                    import json
                    device_ids = json.loads(GPU_DEVICE_IDS)
                    if device_ids and isinstance(device_ids, list):
                        # 使用第一个GPU设备（容器内的设备号）
                        device_id = device_ids[0]
                        device = f"cuda:{device_id}"
                        logger.info(f"使用GPU设备: {device} (从GPU_DEVICE_IDS解析，未设置CUDA_VISIBLE_DEVICES)")
                        return device
                except (json.JSONDecodeError, ValueError, IndexError, TypeError) as e:
                    # 如果解析失败，使用默认cuda:0
                    logger.warning(f"解析GPU_DEVICE_IDS失败: {e}，使用默认cuda:0设备")
                
                # 默认使用第一个GPU（cuda:0）
                device = "cuda:0"
                logger.info(f"使用GPU设备: {device} (默认)")
                return device
            else:
                logger.warning("CUDA不可用，回退到CPU")
                return "cpu"
        except ImportError:
            # PyTorch未安装，无法使用CUDA
            logger.warning("PyTorch未安装，无法使用CUDA，回退到CPU")
            return "cpu"
    else:
        logger.info("GPU配置: USE_GPU=false，使用CPU设备")
        return "cpu"

# 配置日志系统
def setup_logging():
    """配置日志系统，同时输出到控制台和文件"""
    # 创建日志目录
    log_dir = Path("/app/logs")
    log_dir.mkdir(parents=True, exist_ok=True)
    
    # 日志文件名包含端口号，方便区分不同实例
    log_file = log_dir / f"cv_api_port_{PORT}.log"
    
    # 配置日志格式
    log_format = '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    date_format = '%Y-%m-%d %H:%M:%S'
    
    # 配置根日志记录器
    logger = logging.getLogger()
    logger.setLevel(logging.INFO)
    
    # 清除已有的处理器
    logger.handlers.clear()
    
    # 控制台处理器
    console_handler = logging.StreamHandler(sys.stdout)
    console_handler.setLevel(logging.INFO)
    console_formatter = logging.Formatter(log_format, date_format)
    console_handler.setFormatter(console_formatter)
    
    # 文件处理器
    file_handler = logging.FileHandler(log_file, encoding='utf-8', mode='a')
    file_handler.setLevel(logging.INFO)
    file_formatter = logging.Formatter(log_format, date_format)
    file_handler.setFormatter(file_formatter)
    
    # 添加处理器
    logger.addHandler(console_handler)
    logger.addHandler(file_handler)
    
    return logger

# 初始化日志系统
logger = setup_logging()
logger.info("=" * 60)
logger.info(f"CV API 服务启动 - 端口: {PORT}")
logger.info(f"日志文件: /app/logs/cv_api_port_{PORT}.log")
logger.info("=" * 60)

# 检查至少配置了一个模型
if MODEL_PATH is None and SEGMENT_MODEL_PATH is None:
    raise ValueError("至少需要配置一个模型：MODEL_PATH（分类模型）或 SEGMENT_MODEL_PATH（分割模型）")

# API配置
# 根据配置的模型动态生成标题和描述
title_parts = []
if MODEL_PATH:
    title_parts.append("分类")
if SEGMENT_MODEL_PATH:
    title_parts.append("分割")
title = "图像" + "与".join(title_parts) + "服务API" if title_parts else "图像处理服务API"

description_parts = []
if MODEL_PATH:
    description_parts.append("图像分类")
if SEGMENT_MODEL_PATH:
    description_parts.append("图像分割")
description = "基于深度学习的" + "和".join(description_parts) + "API服务" if description_parts else "图像处理API服务"

API_CONFIG = {
    "host": HOST,
    "port": PORT,
    "title": title,
    "description": description,
    "version": "1.0.0"
}


class ClassificationRequest(BaseModel):
    """分类请求模型"""
    image: str  # base64编码的图片字符串


class ClassificationResponse(BaseModel):
    """分类响应模型"""
    success: bool
    classifier_name: str
    class_id: int
    class_name: str
    confidence: float
    probabilities: list
    message: Optional[str] = None


class SegmentationRequest(BaseModel):
    """分割请求模型"""
    image: str  # base64编码的图片字符串


class SegmentationResponse(BaseModel):
    """分割响应模型"""
    success: bool
    segmentation_mask_base64: Optional[str] = None  # base64编码的分割掩码图像
    color_mask_base64: Optional[str] = None  # base64编码的彩色掩码图像
    num_classes: Optional[int] = None
    image_shape: Optional[List[int]] = None
    class_names: Optional[List[str]] = None
    message: Optional[str] = None


# 创建FastAPI应用
app = FastAPI(
    title=API_CONFIG["title"],
    description=API_CONFIG["description"],
    version=API_CONFIG["version"]
)


@app.exception_handler(RequestValidationError)
async def validation_exception_handler(request: Request, exc: RequestValidationError):
    """
    处理请求验证错误，当接收的数据字段不是image时返回友好的中文错误信息
    """
    errors = exc.errors()
    
    # 检查所有错误，判断是否与字段名相关
    has_non_image_field = False
    missing_image_field = False
    
    for error in errors:
        error_loc = error.get('loc', [])
        error_type = error.get('type', '')
        
        # 检查是否是body字段错误
        if len(error_loc) > 1 and error_loc[0] == 'body':
            field_name = error_loc[-1] if error_loc else None
            
            # 如果字段名不是image，说明有额外的字段或错误的字段名
            if field_name != 'image':
                has_non_image_field = True
            # 如果字段名是image但类型是missing，说明缺少image字段
            elif field_name == 'image' and 'missing' in error_type:
                missing_image_field = True
    
    # 如果有非image字段或缺少image字段，返回详细错误信息
    if has_non_image_field or missing_image_field:
        return JSONResponse(
            status_code=400,
            content={"detail": "数据输入有误：请求数据中必须包含image字段"}
        )
    
    # 其他验证错误也返回数据输入有误
    return JSONResponse(
        status_code=400,
        content={"detail": "数据输入有误"}
    )

# 全局分类器实例
classifier = None

# 全局分割器实例
segmenter = None


def get_classifier():
    """获取分类器实例"""
    global classifier
    
    if MODEL_PATH is None:
        raise HTTPException(status_code=503, detail="分类模型未配置，请设置环境变量 MODEL_PATH")
    
    if classifier is None:
        device = get_device()
        logger.info("正在加载分类器...")
        logger.info(f"模型路径: {MODEL_PATH}")
        logger.info(f"设备: {device}")
        logger.info(f"GPU配置: USE_GPU={USE_GPU}, GPU_DEVICE_IDS={GPU_DEVICE_IDS}")
        
        try:
            classifier = Classifier(
                model_path=MODEL_PATH,
                device=device,
                class_names=None,  # 自动从模型配置中获取
                logger=logger  # 传入logger，用于记录设备信息
            )
            logger.info("✅ 成功加载分类器")
        except Exception as e:
            logger.error(f"❌ 加载分类器失败: {e}")
            import traceback
            error_detail = traceback.format_exc()
            logger.error(f"错误详情:\n{error_detail}")
            raise
    
    return classifier


def get_segmenter():
    """获取分割器实例"""
    global segmenter
    
    if SEGMENT_MODEL_PATH is None:
        raise HTTPException(status_code=503, detail="分割模型未配置，请设置环境变量 SEGMENT_MODEL_PATH")
    
    if segmenter is None:
        device = get_device()
        logger.info("正在加载分割器...")
        logger.info(f"模型路径: {SEGMENT_MODEL_PATH}")
        logger.info(f"设备: {device}")
        logger.info(f"GPU配置: USE_GPU={USE_GPU}, GPU_DEVICE_IDS={GPU_DEVICE_IDS}")
        
        try:
            segmenter = Segmenter(
                model_path=SEGMENT_MODEL_PATH,
                device=device,
                class_names=None,  # 自动从模型配置中获取
                logger=logger  # 传入logger，用于记录设备信息
            )
            logger.info("✅ 成功加载分割器")
        except Exception as e:
            logger.error(f"❌ 加载分割器失败: {e}")
            import traceback
            error_detail = traceback.format_exc()
            logger.error(f"错误详情:\n{error_detail}")
            raise
    
    return segmenter


def numpy_to_base64(mask: np.ndarray, format: str = 'PNG') -> str:
    """
    将numpy数组（分割掩码）转换为base64编码的图像字符串
    
    Args:
        mask: 分割掩码numpy数组
        format: 图像格式，默认为PNG
        
    Returns:
        str: base64编码的图像字符串（包含数据URI前缀）
    """
    try:
        # 确保数据类型正确
        if mask.dtype != np.uint8:
            # 如果掩码值在合理范围内，转换为uint8
            if mask.max() <= 255 and mask.min() >= 0:
                mask = mask.astype(np.uint8)
            else:
                # 如果值超出范围，进行归一化
                mask = ((mask - mask.min()) / (mask.max() - mask.min() + 1e-8) * 255).astype(np.uint8)
        
        # 将numpy数组转换为PIL Image
        if len(mask.shape) == 2:
            # 单通道灰度图
            pil_image = Image.fromarray(mask, mode='L')
        elif len(mask.shape) == 3:
            # 多通道图像
            pil_image = Image.fromarray(mask)
        else:
            raise ValueError(f"不支持的掩码维度: {mask.shape}")
        
        # 将PIL Image转换为字节流
        buffer = io.BytesIO()
        pil_image.save(buffer, format=format)
        buffer.seek(0)
        
        # 编码为base64
        image_bytes = buffer.read()
        image_base64 = base64.b64encode(image_bytes).decode('utf-8')
        
        # 返回带数据URI前缀的base64字符串
        mime_type = 'image/png' if format == 'PNG' else 'image/jpeg'
        return f"data:{mime_type};base64,{image_base64}"
        
    except Exception as e:
        raise ValueError(f"转换numpy数组到base64失败: {str(e)}")


def generate_color_mask(mask_array: np.ndarray) -> Image.Image:
    """
    将分割掩码转换为彩色显示（应用颜色映射）
    
    基于view_seg_result.py的逻辑，为每个类别生成随机颜色
    
    Args:
        mask_array: 分割掩码numpy数组（2D或3D）
        
    Returns:
        PIL.Image: 彩色掩码图像（RGB格式）
    """
    import random
    
    # 确保是numpy数组
    if not isinstance(mask_array, np.ndarray):
        mask_array = np.array(mask_array)
    
    # 如果是3D数组，先转换为2D
    if len(mask_array.shape) == 3:
        # 如果是多通道，取第一个通道或转换为灰度
        if mask_array.shape[2] == 1:
            mask_array = mask_array[:, :, 0]
        else:
            # 多通道情况，可能需要特殊处理，这里先转换为灰度
            mask_array = np.mean(mask_array, axis=2)
    
    # 确保数据类型正确：转换为uint8
    # 处理可能的负数或超出范围的值
    if mask_array.dtype != np.uint8:
        # 如果有负数，先偏移到0
        if mask_array.min() < 0:
            mask_array = mask_array - mask_array.min()
        # 确保值在合理范围内（0-255）
        if mask_array.max() > 255:
            # 如果最大值超过255，进行归一化
            mask_array = ((mask_array - mask_array.min()) / (mask_array.max() - mask_array.min() + 1e-8) * 255).astype(np.uint8)
        else:
            mask_array = mask_array.astype(np.uint8)
    
    if len(mask_array.shape) == 2:
        # 基于mask中的实际类别数量自动生成调色板
        num_classes = int(mask_array.max()) + 1
        palette = []
        # 0类为黑色
        palette.extend([0, 0, 0])
        # 其他类别随机生成颜色
        for cls in range(1, num_classes):
            color = [random.randint(0, 255) for _ in range(3)]
            palette.extend(color)
        # 需要填充到256*3长度
        palette += [0, 0, 0] * (256 - num_classes)
        colored = Image.fromarray(mask_array, mode='L').convert('P')
        colored.putpalette(palette)
        colored = colored.convert('RGB')
    else:
        # 如果已经是RGB图像，直接返回
        colored = Image.fromarray(mask_array)
    
    return colored


@app.get("/")
async def root():
    """根路径"""
    available_services = []
    if MODEL_PATH:
        available_services.append("分类服务")
    if SEGMENT_MODEL_PATH:
        available_services.append("分割服务")
    
    return {
        "message": "图像API服务正在运行",
        "available_services": available_services,
        "classification": MODEL_PATH is not None,
        "segmentation": SEGMENT_MODEL_PATH is not None
    }


@app.get("/health")
async def health_check():
    """健康检查"""
    available_services = []
    if MODEL_PATH:
        available_services.append("分类服务")
    if SEGMENT_MODEL_PATH:
        available_services.append("分割服务")
    
    return {
        "status": "healthy",
        "message": "API服务正常运行",
        "available_services": available_services
    }


@app.get("/config")
async def get_config():
    """获取配置信息"""
    return {
        "classification_model_path": MODEL_PATH,
        "segmentation_model_path": SEGMENT_MODEL_PATH,
        "host": HOST,
        "port": PORT,
        "classification_enabled": MODEL_PATH is not None,
        "segmentation_enabled": SEGMENT_MODEL_PATH is not None
    }


@app.post("/api/v1/classify/single", response_model=ClassificationResponse)
async def classify_single_image(request: ClassificationRequest):
    """
    分类单张图片（使用base64编码）
    
    Args:
        request: 包含base64编码图片的请求对象
        
    Returns:
        ClassificationResponse: 分类结果
    """
    try:
        # 解码base64图片
        image_base64 = request.image
        
        # 处理base64字符串，移除可能的数据URI前缀（如 "data:image/jpeg;base64,"）
        if ',' in image_base64:
            image_base64 = image_base64.split(',')[1]
        
        # 解码base64数据
        try:
            image_data = base64.b64decode(image_base64)
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"Base64解码失败: {str(e)}")
        
        # 将字节数据转换为PIL Image
        try:
            pil_image = Image.open(io.BytesIO(image_data))
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"图片格式错误: {str(e)}")
        
        # 获取分类器
        classifier = get_classifier()
        
        # 执行分类
        result = classifier.classify_single(pil_image)
        
        if not result.get('success', False):
            raise HTTPException(status_code=400, detail=result.get('error', '分类失败'))
        
        return ClassificationResponse(
            success=True,
            classifier_name="default",
            class_id=int(result['class_id']),
            class_name=result['class_name'],
            confidence=float(result['confidence']),
            probabilities=[float(p) for p in result.get('probabilities', [])],
            message=result.get('message')
        )
        
    except HTTPException:
        raise
    except Exception as e:
        import traceback
        error_detail = traceback.format_exc()
        logger.error(f"分类API错误详情:\n{error_detail}")
        raise HTTPException(status_code=500, detail=f"分类过程中发生错误: {str(e)}")


@app.post("/api/v1/segment/single", response_model=SegmentationResponse)
async def segment_single_image(request: SegmentationRequest):
    """
    分割单张图片（使用base64编码）
    
    Args:
        request: 包含base64编码图片的请求对象
        
    Returns:
        SegmentationResponse: 分割结果，包含base64编码的分割掩码图像
    """
    try:
        # 解码base64图片
        image_base64 = request.image
        
        # 处理base64字符串，移除可能的数据URI前缀（如 "data:image/jpeg;base64,"）
        if ',' in image_base64:
            image_base64 = image_base64.split(',')[1]
        
        # 解码base64数据
        try:
            image_data = base64.b64decode(image_base64)
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"Base64解码失败: {str(e)}")
        
        # 将字节数据转换为PIL Image
        try:
            pil_image = Image.open(io.BytesIO(image_data))
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"图片格式错误: {str(e)}")
        
        # 保存原始图像的副本和尺寸（用于后续叠加）
        original_image = pil_image.copy()
        original_size = original_image.size
        
        # 获取分割器
        segmenter = get_segmenter()
        
        # 执行分割
        result = segmenter.segment_single(pil_image)
        
        if not result.get('success', False):
            raise HTTPException(status_code=400, detail=result.get('error', '分割失败'))
        
        # 获取分割掩码
        segmentation_mask = result.get('segmentation_mask')
        if segmentation_mask is None:
            raise HTTPException(status_code=500, detail="分割结果中没有分割掩码")
        
        # 将numpy数组转换为base64编码的图像
        try:
            mask_base64 = numpy_to_base64(segmentation_mask, format='PNG')
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"转换分割掩码为base64失败: {str(e)}")
        
        # 生成彩色掩码并叠加到原图上，然后转换为base64
        try:
            # 生成彩色掩码
            color_mask_image = generate_color_mask(segmentation_mask)
            logger.info(f"彩色掩码生成成功，尺寸: {color_mask_image.size}, 模式: {color_mask_image.mode}")
            
            # 确保原图是RGB格式
            if original_image.mode != 'RGB':
                original_image = original_image.convert('RGB')
            
            # 获取分割掩码的尺寸（可能是模型输出尺寸，不是原图尺寸）
            mask_height, mask_width = segmentation_mask.shape[:2]
            mask_size = (mask_width, mask_height)  # PIL使用(width, height)
            
            logger.info(f"原图尺寸: {original_size}, 掩码尺寸: {mask_size}, 彩色掩码尺寸: {color_mask_image.size}")
            
            # 如果彩色掩码尺寸和原图不一致，需要先resize彩色掩码到原图尺寸
            if color_mask_image.size != original_size:
                logger.info(f"调整彩色掩码尺寸从 {color_mask_image.size} 到 {original_size}")
                color_mask_image = color_mask_image.resize(original_size, Image.Resampling.LANCZOS)
            
            # 确保两个图像尺寸完全一致
            if color_mask_image.size != original_image.size:
                logger.warning(f"尺寸不匹配: 原图 {original_image.size}, 彩色掩码 {color_mask_image.size}, 强制调整")
                color_mask_image = color_mask_image.resize(original_image.size, Image.Resampling.LANCZOS)
            
            # 确保两个图像都是RGB模式
            if color_mask_image.mode != 'RGB':
                color_mask_image = color_mask_image.convert('RGB')
            
            # 将彩色掩码叠加到原图上（透明度0.5，可以同时看到原图和掩码）
            overlay_image = Image.blend(original_image, color_mask_image, alpha=0.5)
            
            # 将叠加后的图像转换为base64
            buffer = io.BytesIO()
            overlay_image.save(buffer, format='PNG')
            buffer.seek(0)
            color_mask_bytes = buffer.read()
            color_mask_base64 = base64.b64encode(color_mask_bytes).decode('utf-8')
            color_mask_base64 = f"data:image/png;base64,{color_mask_base64}"
            logger.info("彩色掩码叠加图生成成功")
        except Exception as e:
            import traceback
            error_detail = traceback.format_exc()
            logger.error(f"生成彩色掩码叠加图失败: {str(e)}")
            logger.error(f"错误详情:\n{error_detail}")
            color_mask_base64 = None
        
        # 转换image_shape为列表（numpy数组的shape是tuple，需要转换为list以便JSON序列化）
        image_shape = None
        if 'image_shape' in result and result['image_shape'] is not None:
            image_shape = list(result['image_shape'])
        
        return SegmentationResponse(
            success=True,
            segmentation_mask_base64=mask_base64,
            color_mask_base64=color_mask_base64,
            num_classes=result.get('num_classes'),
            image_shape=image_shape,
            class_names=result.get('class_names'),
            message=result.get('message')
        )
        
    except HTTPException:
        raise
    except Exception as e:
        import traceback
        error_detail = traceback.format_exc()
        logger.error(f"分割API错误详情:\n{error_detail}")
        raise HTTPException(status_code=500, detail=f"分割过程中发生错误: {str(e)}")


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host=API_CONFIG["host"], port=API_CONFIG["port"])
